import pandas as pd
from scipy.stats import norm

from ..multiprocessing import mp_pandas_obj


def get_signal(events: pd.DataFrame,
               step_size,
               prob: pd.Series,
               pred: pd.Series,
               num_cls, num_threads, **kwargs):
    """
    :param events:
    :param step_size:
    :param prob: prob的index要唯一，且包含在events中
    :param pred:
    :param num_cls:
    :param num_threads:
    :param kwargs:
    :return:
    """
    # get signals from predictions
    if prob.shape[0] == 0:
        return pd.Series
    # 1) generate signals from multinomial classification (one-vs-rest, OvR)
    signal0: pd.Series = (prob - 1. / num_cls) / (prob * (1. - prob)) ** .5  # t-value of OvR
    signal0: pd.Series = pred * (2 * norm.cdf(signal0) - 1)  # signal=side*size
    if "side" in events:
        signal0 *= events.loc[signal0.index, "side"]  # meta-labelling
        # 2) compute average signal among those concurrently open
    # 2) compute average signal among those concurrently open
    df = signal0.to_frame("signal").join(events[["t1"]], how="left")
    df = avg_active_signals(df, num_threads)
    signal1 = discrete_signal(signal0=df, step_size=step_size)
    return signal1


def mp_avg_active_signals(signals, molecule):
    out = pd.Series()
    for loc in molecule:
        df0 = (signals.index.values <= loc) & ((loc < signals['t1']) | pd.isnull(signals['t1']))
        act = signals[df0].index
        if len(act) > 0:
            out[loc] = signals.loc[act, 'signal'].mean()
        else:
            out[loc] = 0  # no signals active at this time
    return out


def avg_active_signals(signals, numThreads):
    # compute the average signal among those active
    # 1) time points where signals change (either one starts or one ends)
    tPnts = set(signals['t1'].dropna().values)
    tPnts = tPnts.union(signals.index.values)
    tPnts = list(tPnts)
    tPnts.sort()
    out = mp_pandas_obj(mp_avg_active_signals, ('molecule', tPnts), numThreads, signals=signals)
    return out


def discrete_signal(signal0, step_size):
    """

    I suggest you discretize the bet size as m ∗= round[ m/d ]d,
    where d ∈ (0, 1] determines the degree of discretization

    :param signal0:
    :param step_size:
    :return:
    """
    # discretize signal
    signal1 = (signal0 / step_size).round() * step_size  # discretize
    signal1[signal1 > 1] = 1  # cap
    signal1[signal1 < -1] = -1  # floor
    return signal1
